# asm_to_kg.py
import re
import networkx as nx
from collections import defaultdict

def extract_functions_and_calls(asm_content):
    """
    从汇编代码中提取函数和调用关系
    """
    functions = {}
    calls = defaultdict(list)
    current_function = None
    function_body = []
    
    # 正则表达式匹配函数标签和调用指令
    function_pattern = re.compile(r'^(\.text:\d+\s+)?([a-zA-Z_][a-zA-Z0-9_]*)(?::|\s+proc\s+near)')
    call_pattern = re.compile(r'\bcall\b\s+(.+)')

    for line in asm_content.split('\n'):
        # 检查函数定义
        func_match = function_pattern.match(line)
        if func_match:
            if current_function:
                # 保存前一个函数体
                functions[current_function] = function_body
                function_body = []
            current_function = func_match.group(2)
            continue
        
        # 检查调用指令
        if current_function:
            call_match = call_pattern.search(line)
            if call_match:
                callee = call_match.group(1)
                calls[current_function].append(callee)
            # 将指令添加到当前函数体
            function_body.append(line.strip())
    
    # 保存最后一个函数
    if current_function:
        functions[current_function] = function_body
    
    return functions, calls

def build_knowledge_graph(functions, calls):
    """
    构建知识图谱
    """
    G = nx.DiGraph()
    
    # 添加函数节点
    for func in functions:
        # 使用函数体前10条指令作为节点特征
        body_snippet = " ".join(functions[func][:10])
        G.add_node(func, body=body_snippet)
    
    # 添加调用关系边
    for caller, callees in calls.items():
        for callee in callees:
            if callee in functions:  # 确保被调用函数存在
                G.add_edge(caller, callee)
    
    return G

def graph_to_text(graph):
    """
    将知识图谱转换为文本描述
    """
    text_lines = []
    
    # 添加节点描述
    for node in graph.nodes:
        body_snippet = graph.nodes[node].get('body', '')
        text_lines.append(f"Function {node} contains instructions: {body_snippet}")
    
    # 添加边描述
    for src, dst in graph.edges:
        text_lines.append(f"Function {src} calls function {dst}")
    
    return "\n".join(text_lines)

# 接口: 处理汇编文件并生成知识图谱文本
def process_asm_to_graph_text(asm_path):
    """
    处理汇编文件并生成知识图谱文本表示

    Args:
        asm_path: 汇编文件路径
    
    Returns:
        graph_text: 知识图谱的文本表示
    """
    with open(asm_path, 'r', encoding='utf-8', errors='ignore') as f:
        asm_content = f.read()

    # 提取函数和调用关系
    functions, calls = extract_functions_and_calls(asm_content)
    
    # 构建知识图谱
    graph = build_knowledge_graph(functions, calls)
    
    # 将图谱转换为文本
    graph_text = graph_to_text(graph)
    
    return graph_text